# Copyright (c) Facebook, Inc. and its affiliates. All Rights Reserved
"""
DETR Transformer class.

Copy-paste from torch.nn.Transformer with modifications:
    * positional encodings are passed in MHattention
    * extra LN at the end of encoder is removed
    * decoder returns a stack of activations from all decoding layers
"""
import copy
from typing import Optional, List

import torch
import torch.nn.functional as F
from torch import nn, Tensor

"""参数：
d_model：模型的嵌入维度，默认为 512。
nhead：多头注意力机制的头数，默认为 8。
num_encoder_layers：编码器的层数，默认为 6。
num_decoder_layers：解码器的层数，默认为 6。
dim_feedforward：前馈神经网络的隐藏层维度，默认为 2048。
dropout：Dropout 概率，默认为 0.1。
activation：激活函数，默认为 "relu"。
normalize_before：是否在子层之前进行 LayerNorm，默认为 False。
return_intermediate_dec：是否返回解码器每一层的输出，默认为 False。"""
class Transformer(nn.Module):#初始化transformer

    def __init__(self, d_model=512, nhead=8, num_encoder_layers=6,
                 num_decoder_layers=6, dim_feedforward=2048, dropout=0.1,
                 activation="relu", normalize_before=False,
                 return_intermediate_dec=False):
        super().__init__()

        #encoder_layer：编码器的单层实例，由 TransformerEncoderLayer 实现。
        encoder_layer = TransformerEncoderLayer(d_model, nhead, dim_feedforward,
                                                dropout, activation, normalize_before)
        #encoder_norm：如果 normalize_before 为 True，则在编码器末尾添加 LayerNorm；否则为 None。
        encoder_norm = nn.LayerNorm(d_model) if normalize_before else None
        #self.encoder：将单层编码器堆叠成多层编码器，由 TransformerEncoder 实现。
        self.encoder = TransformerEncoder(encoder_layer, num_encoder_layers, encoder_norm)

        #decoder_layer：解码器的单层实例，由 TransformerDecoderLayer 实现。
        decoder_layer = TransformerDecoderLayer(d_model, nhead, dim_feedforward,
                                                dropout, activation, normalize_before)
        #decoder_norm：在解码器末尾添加 LayerNorm。
        decoder_norm = nn.LayerNorm(d_model)
        #elf.decoder：将单层解码器堆叠成多层解码器，由 TransformerDecoder 实现。
        #return_intermediate：如果为 True，则返回解码器每一层的输出。
        self.decoder = TransformerDecoder(decoder_layer, num_decoder_layers, decoder_norm,
                                          return_intermediate=return_intermediate_dec) 

        self._reset_parameters()

        self.d_model = d_model
        self.nhead = nhead

    def _reset_parameters(self):#定义 _reset_parameters 方法，用于初始化模型参数
        for p in self.parameters():#遍历模型的所有参数。
            if p.dim() > 1:#初始化模型参数：对维度大于 1 的参数使用 Xavier 均匀初始化
                nn.init.xavier_uniform_(p)

    def forward(self, src, mask, query_embed, pos_embed):
        # flatten NxCxHxW to HWxNxC
        bs, c, h, w = src.shape
        #第 2 维开始（包括第 2 维）的所有维度展平为一维,输入形状是 (bs, c, h, w)展平后，形状变为 (bs, c, h * w)
        src = src.flatten(2).permute(2, 0, 1)#将图像特征从 [bs, c, h, w] 展平为 [h*w, bs, c]
        pos_embed = pos_embed.flatten(2).permute(2, 0, 1)#位置编码展平为 [h*w, bs, c]
        #query_embed**：是可学习的参数（nn.Embedding(num_queries, d_model)），表示 object queries 的初始嵌入。
        query_embed = query_embed.unsqueeze(1).repeat(1, bs, 1)#扩展为 [num_queries, bs, c]
        """ repeat 的作用是沿着指定的维度复制张量。在这里，repeat(1, bs, 1) 的含义是：沿着第 0 维（n_queries）不复
        制（保持不变）。沿着第 1 维（新增的批量维度）复制 bs 次。沿着第 2 维（embed_dim）不复制（保持不变）"""
        mask = mask.flatten(1)#展平为 [bs, h*w]

        """在 DETR 的 Transformer 实现中，tgt 确实代表 目标序列（target sequence），也就是 可学
        习的 object queries.
        在 Transformer.forward() 中，tgt 被初始化为全零张量，形状与 query_embed 相同：
        tgt 是解码器的输入序列，每个 query 对应一个潜在的物体预测（如 100 个 queries 对应 COCO 数据
        集的 100 个物体）"""

        tgt = torch.zeros_like(query_embed)#初始化为全零，形状与 query_embed 相同。
        memory = self.encoder(src, src_key_padding_mask=mask, pos=pos_embed)#编码器输出
        hs = self.decoder(tgt, memory, memory_key_padding_mask=mask,#解码器输出
                          pos=pos_embed, query_pos=query_embed)
        return hs.transpose(1, 2), memory.permute(1, 2, 0).view(bs, c, h, w)# 编码器输出[h*w, bs, c]变为(bs,c,h,w)
"""Q: 为什么需要 query_embed 和 pos_embed 两种位置编码？
    query_embed：为 object queries 提供初始位置信息（可学习）。
    pos_embed：为图像特征提供空间位置信息（固定或可学习）。
    两者共同帮助模型建立物体与图像位置的关联。
tgt 如何从零变为最终的物体预测？通过解码器多层堆叠的注意力机制，逐步融合图像特征（memory）和 queries 的位
置信息（query_embed），最终输出物体特征。"""

class TransformerEncoder(nn.Module):#初始化编码器：encoder_layer：单层编码器。num_layers：编码器层数。norm：LayerNorm 层。

    def __init__(self, encoder_layer, num_layers, norm=None):
        super().__init__()
        self.layers = _get_clones(encoder_layer, num_layers)
        self.num_layers = num_layers
        self.norm = norm

    def forward(self, src,
                mask: Optional[Tensor] = None,
                src_key_padding_mask: Optional[Tensor] = None,
                pos: Optional[Tensor] = None):
        output = src#前向传播：遍历所有编码器层，逐层处理输入。

        for layer in self.layers:
            output = layer(output, src_mask=mask,
                           src_key_padding_mask=src_key_padding_mask, pos=pos)

        if self.norm is not None:#如果 norm 不为 None，对输出进行 LayerNorm。
            output = self.norm(output)

        return output


class TransformerDecoder(nn.Module):
    #初始化解码器：decoder_layer：单层解码器。num_layers：解码器层数。norm：LayerNorm 层。return_intermediate：是否返回中间层输出。
    def __init__(self, decoder_layer, num_layers, norm=None, return_intermediate=False):
        super().__init__()
        self.layers = _get_clones(decoder_layer, num_layers)
        self.num_layers = num_layers
        self.norm = norm
        self.return_intermediate = return_intermediate

    def forward(self, tgt, memory,
                tgt_mask: Optional[Tensor] = None,
                memory_mask: Optional[Tensor] = None,
                tgt_key_padding_mask: Optional[Tensor] = None,
                memory_key_padding_mask: Optional[Tensor] = None,
                pos: Optional[Tensor] = None,
                query_pos: Optional[Tensor] = None):
        output = tgt

        intermediate = []

        for layer in self.layers:#遍历所有解码器层，逐层处理输入。
            output = layer(output, memory, tgt_mask=tgt_mask,
                           memory_mask=memory_mask,
                           tgt_key_padding_mask=tgt_key_padding_mask,
                           memory_key_padding_mask=memory_key_padding_mask,
                           pos=pos, query_pos=query_pos)
            if self.return_intermediate:#如果 return_intermediate 为 True，保存中间层输出。
                intermediate.append(self.norm(output))

        if self.norm is not None:#对输出进行 LayerNorm。
            output = self.norm(output)
            if self.return_intermediate:
                intermediate.pop()
                intermediate.append(output)

        if self.return_intermediate:#返回中间层输出或最终输出。
            return torch.stack(intermediate)

        return output.unsqueeze(0)


class TransformerEncoderLayer(nn.Module):
    #初始化编码器层：自注意力层 self_attn。前馈网络 linear1 和 linear2。LayerNorm 层 norm1 和 norm2。Dropout 层。激活函数 activation。
    #normalize_before：是否在子层前进行 LayerNorm。
    def __init__(self, d_model, nhead, dim_feedforward=2048, dropout=0.1,
                 activation="relu", normalize_before=False):
        super().__init__()
        self.self_attn = nn.MultiheadAttention(d_model, nhead, dropout=dropout)
        # Implementation of Feedforward model
        self.linear1 = nn.Linear(d_model, dim_feedforward)
        self.dropout = nn.Dropout(dropout)
        self.linear2 = nn.Linear(dim_feedforward, d_model)

        self.norm1 = nn.LayerNorm(d_model)
        self.norm2 = nn.LayerNorm(d_model)
        self.dropout1 = nn.Dropout(dropout)
        self.dropout2 = nn.Dropout(dropout)

        self.activation = _get_activation_fn(activation)
        self.normalize_before = normalize_before

    def with_pos_embed(self, tensor, pos: Optional[Tensor]):
        return tensor if pos is None else tensor + pos

    def forward_post(self,#后归一化前向传播：
                     src,#src是从backbone的图片，850*bs*256
                     src_mask: Optional[Tensor] = None,
                     src_key_padding_mask: Optional[Tensor] = None,
                     pos: Optional[Tensor] = None):
        #在 TransformerDecoderLayer 中，query_embed 作为位置编码加到 tgt
        q = k = self.with_pos_embed(src, pos)#src+pos得到q和k
        src2 = self.self_attn(q, k, value=src, attn_mask=src_mask,#自注意力计算。
                              key_padding_mask=src_key_padding_mask)[0]
        src = src + self.dropout1(src2)#残差连接
        src = self.norm1(src)#layernorm
        src2 = self.linear2(self.dropout(self.activation(self.linear1(src))))#FFN 前馈网络计算
        src = src + self.dropout2(src2)
        src = self.norm2(src)#encoder的第二个layernorm
        return src

    def forward_pre(self, src,#前归一化前向传播 应该是norm的位置不同
                    src_mask: Optional[Tensor] = None,
                    src_key_padding_mask: Optional[Tensor] = None,
                    pos: Optional[Tensor] = None):
        src2 = self.norm1(src)#先进行 LayerNorm
        q = k = self.with_pos_embed(src2, pos)
        src2 = self.self_attn(q, k, value=src2, attn_mask=src_mask,#自注意力计算
                              key_padding_mask=src_key_padding_mask)[0]
        src = src + self.dropout1(src2)#残差连接。
        src2 = self.norm2(src)#
        src2 = self.linear2(self.dropout(self.activation(self.linear1(src2))))#前馈网络计算。
        src = src + self.dropout2(src2)
        return src

    def forward(self, src,#根据 normalize_before 选择前向传播方式
                src_mask: Optional[Tensor] = None,
                src_key_padding_mask: Optional[Tensor] = None,
                pos: Optional[Tensor] = None):
        if self.normalize_before:
            return self.forward_pre(src, src_mask, src_key_padding_mask, pos)
        return self.forward_post(src, src_mask, src_key_padding_mask, pos)


class TransformerDecoderLayer(nn.Module):
    #初始化解码器层：自注意力层 self_attn。交叉注意力层 multihead_attn。前馈网络 linear1 和 linear2。
    # LayerNorm 层 norm1、norm2 和 norm3。Dropout 层。激活函数 activation。normalize_before：是否在子层前进行 LayerNorm。
    def __init__(self, d_model, nhead, dim_feedforward=2048, dropout=0.1,
                 activation="relu", normalize_before=False):
        super().__init__()
        self.self_attn = nn.MultiheadAttention(d_model, nhead, dropout=dropout)
        self.multihead_attn = nn.MultiheadAttention(d_model, nhead, dropout=dropout)
        # Implementation of Feedforward model
        self.linear1 = nn.Linear(d_model, dim_feedforward)
        self.dropout = nn.Dropout(dropout)
        self.linear2 = nn.Linear(dim_feedforward, d_model)

        self.norm1 = nn.LayerNorm(d_model)
        self.norm2 = nn.LayerNorm(d_model)
        self.norm3 = nn.LayerNorm(d_model)
        self.dropout1 = nn.Dropout(dropout)
        self.dropout2 = nn.Dropout(dropout)
        self.dropout3 = nn.Dropout(dropout)

        self.activation = _get_activation_fn(activation)
        self.normalize_before = normalize_before

    def with_pos_embed(self, tensor, pos: Optional[Tensor]):
        return tensor if pos is None else tensor + pos

    def forward_post(self, tgt, memory,#后归一化前向传播： 
                     tgt_mask: Optional[Tensor] = None,#tgt就是queries，默认初始值0，(100,256)
                     memory_mask: Optional[Tensor] = None,#enceder的输出
                     tgt_key_padding_mask: Optional[Tensor] = None,
                     memory_key_padding_mask: Optional[Tensor] = None,
                     pos: Optional[Tensor] = None,
                     query_pos: Optional[Tensor] = None):#object queries尺寸：(100,256)
        q = k = self.with_pos_embed(tgt, query_pos)#相加得到q和k
        tgt2 = self.self_attn(q, k, value=tgt, attn_mask=tgt_mask,#qkv进行attn操作
                              key_padding_mask=tgt_key_padding_mask)[0]
        tgt = tgt + self.dropout1(tgt2)
        tgt = self.norm1(tgt)
        tgt2 = self.multihead_attn(query=self.with_pos_embed(tgt, query_pos),
                                   key=self.with_pos_embed(memory, pos),
                                   value=memory, attn_mask=memory_mask,
                                   key_padding_mask=memory_key_padding_mask)[0]
        tgt = tgt + self.dropout2(tgt2)
        tgt = self.norm2(tgt)
        tgt2 = self.linear2(self.dropout(self.activation(self.linear1(tgt))))#FFN层 前馈网络计算。
        tgt = tgt + self.dropout3(tgt2)
        tgt = self.norm3(tgt)#输出val (bs,100,256) train(6,bs,100,256)训练记录下了中间值
        return tgt

    def forward_pre(self, tgt, memory,#前归一化前向传播：
                    tgt_mask: Optional[Tensor] = None,
                    memory_mask: Optional[Tensor] = None,
                    tgt_key_padding_mask: Optional[Tensor] = None,
                    memory_key_padding_mask: Optional[Tensor] = None,
                    pos: Optional[Tensor] = None,
                    query_pos: Optional[Tensor] = None):
        tgt2 = self.norm1(tgt)#先进行 LayerNorm。
        q = k = self.with_pos_embed(tgt2, query_pos)#
        tgt2 = self.self_attn(q, k, value=tgt2, attn_mask=tgt_mask,#自注意力计算。
                              key_padding_mask=tgt_key_padding_mask)[0]
        tgt = tgt + self.dropout1(tgt2)#残差连接。
        tgt2 = self.norm2(tgt)#
        tgt2 = self.multihead_attn(query=self.with_pos_embed(tgt2, query_pos),#交叉注意力计算。
                                   key=self.with_pos_embed(memory, pos),
                                   value=memory, attn_mask=memory_mask,
                                   key_padding_mask=memory_key_padding_mask)[0]
        tgt = tgt + self.dropout2(tgt2)#残差连接。
        tgt2 = self.norm3(tgt)
        tgt2 = self.linear2(self.dropout(self.activation(self.linear1(tgt2))))#前馈网络计算。
        tgt = tgt + self.dropout3(tgt2)#残差连接。
        return tgt

    def forward(self, tgt, memory,#根据 normalize_before 选择前向传播方式。
                tgt_mask: Optional[Tensor] = None,
                memory_mask: Optional[Tensor] = None,
                tgt_key_padding_mask: Optional[Tensor] = None,
                memory_key_padding_mask: Optional[Tensor] = None,
                pos: Optional[Tensor] = None,
                query_pos: Optional[Tensor] = None):
        if self.normalize_before:
            return self.forward_pre(tgt, memory, tgt_mask, memory_mask,
                                    tgt_key_padding_mask, memory_key_padding_mask, pos, query_pos)
        return self.forward_post(tgt, memory, tgt_mask, memory_mask,
                                 tgt_key_padding_mask, memory_key_padding_mask, pos, query_pos)


def _get_clones(module, N):#克隆模块 N 次，返回 ModuleList。
    return nn.ModuleList([copy.deepcopy(module) for i in range(N)])


def build_transformer(args):#根据 args 构建 Transformer 实例。
    return Transformer(
        d_model=args.hidden_dim,
        dropout=args.dropout,
        nhead=args.nheads,
        dim_feedforward=args.dim_feedforward,
        num_encoder_layers=args.enc_layers,
        num_decoder_layers=args.dec_layers,
        normalize_before=args.pre_norm,
        return_intermediate_dec=True,
    )


def _get_activation_fn(activation):#根据字符串返回激活函数。
    """Return an activation function given a string"""
    if activation == "relu":
        return F.relu
    if activation == "gelu":
        return F.gelu
    if activation == "glu":
        return F.glu
    raise RuntimeError(F"activation should be relu/gelu, not {activation}.")
